#include <map>
#include <set>
#include <string>
#include <vector>
#include <sstream>
#include <iostream>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#include <DRAMImageStream.hpp>
#include <memline.hpp>
#include <HashFunctions.hpp>
#include <PrettyPrints.hpp>
#include <randvals.hpp>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#include <stdlib.h>
#include <assert.h>
#include <pthread.h>
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
using namespace std;
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
#define ASSERT( x )                   // <--- optimized version
// #define ASSERT( x ) assert( x )    // <---     debug version
#define WARN( condition, message )  if( !( condition ) ) { cout << "... WARNING: condition ( " #condition " ) on line " << __LINE__ << " of file " << __FILE__ << " is not met: " << message << endl; }
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
struct line_compare {
    bool operator() ( const line_t * const & lhs, const line_t * const & rhs ) const {
        if( lhs->d[0]  < rhs->d[0] ) { return true; }
        if( lhs->d[0] == rhs->d[0] ) {
           if( lhs->d[1]  < rhs->d[1] ) { return true; }
           if( lhs->d[1] == rhs->d[1] ) {
              if( lhs->d[2]  < rhs->d[2] ) { return true; }
              if( lhs->d[2] == rhs->d[2] ) {
                 if( lhs->d[3]  < rhs->d[3] ) { return true; }
                 if( lhs->d[3] == rhs->d[3] ) {
                    if( lhs->d[4]  < rhs->d[4] ) { return true; }
                    if( lhs->d[4] == rhs->d[4] ) {
                       if( lhs->d[5]  < rhs->d[5] ) { return true; }
                       if( lhs->d[5] == rhs->d[5] ) {
                          if( lhs->d[6]  < rhs->d[6] ) { return true; }
                          if( lhs->d[6] == rhs->d[6] ) {
                             if( lhs->d[7]  < rhs->d[7] ) {
                                return true;
                             }
                          }
                       }
                    }
                 }
              }
           }
           return false;
        }        
        return false;
    }
};
struct page_compare {
    bool operator() ( const page_t * const & lhs, const page_t * const & rhs ) const {

        ASSERT( lhs->NumLines() == rhs->NumLines() );
        const uint32_t wpp = lhs->NumWords();

        const uint64_t * lp = lhs->RawPtr();
        const uint64_t * rp = rhs->RawPtr();
        
        for( uint32_t wordIdx = 0; wordIdx < wpp; wordIdx++ ) {
            if( *lp < *rp ) { return true;  }
            if( *lp > *rp ) { return false; }
            lp++;
            rp++;
        }
        return false;
    }
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
class DedupCounts {
public:
    uint64_t total;
    uint64_t zero;
    uint64_t unique;
    uint64_t rc_tot;
    DedupCounts() {
        total  = 0ULL;
        zero   = 0ULL;
        unique = 0ULL;
        rc_tot = 0ULL;
    }
    void ResetNonUnique( void ) {
        total = 0ULL;
        zero  = 0ULL;
    }
};
class DedupCountByAlloc {
public:
    DedupCounts allmem;
    DedupCounts kernel;
    DedupCounts user;
    void ResetNonUnique( void ) {
        allmem.ResetNonUnique();
        kernel.ResetNonUnique();
        user.ResetNonUnique();
    }
    inline void IncrementTotal   ( const uint64_t allmemPageRC, const uint64_t kernelPageRC, const uint64_t userPageRC  ) {
        allmem.total += allmemPageRC;
        kernel.total += kernelPageRC;
        user.total   += userPageRC;
    }
    inline void IncrementZero    ( const uint64_t allmemPageRC, const uint64_t kernelPageRC, const uint64_t userPageRC  ) {
        allmem.zero += allmemPageRC;
        kernel.zero += kernelPageRC;
        user.zero   += userPageRC;
    }
    inline void IncrementUnique  ( const uint64_t kernelPageRC, const uint64_t userPageRC  ) {
        allmem.unique++;
        if( kernelPageRC ) {
            kernel.unique++;
        }
        if( userPageRC ) {
            user.unique++;
        }
    }
    inline void IncrementRCAccum ( const uint64_t allmemPageRC, const uint64_t kernelPageRC, const uint64_t userPageRC  ) {
        allmem.rc_tot += allmemPageRC;
        kernel.rc_tot += kernelPageRC;
        user.rc_tot   += userPageRC;
    }
};
class AKURefCount {
    // AKURefCount stores three reference counts: "all", "kern", and "user"
    // all:   all of physical memory
    // kern:  memory considered allocated, i.e. by the kernel, includes i/o buffers (e.g.)
    // user:  memory that is allocated as virtual memory, i.e. given to user space
    // ... all contains kernel, and kernel contains user
public:
    uint32_t all;
    uint32_t kern;
    uint32_t user;
    AKURefCount( void ) {
        all  = 0;
        kern = 0;
        user = 0;
    }
    inline void Increment( const uint32_t kernelPageRC, const uint32_t userPageRC  ) {
        all++;
        if( kernelPageRC ) { kern++; }
        if(   userPageRC ) { user++; }
        // -- TODO -- go back to no if/then
        // kern += kernelPageRC;
        // user +=   userPageRC;
    }
    inline bool operator < ( const AKURefCount & other ) const {
        if( all < other.all ) return true;
        if( all > other.all ) return false;
        if( kern < other.kern ) return true;
        if( kern > other.kern ) return false;
        if( user < other.user ) return true;
        return false;
    }
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
class ContentGenerator {
    static const uint32_t hugeWordIdx = 0;
    static const uint32_t pageWordIdx = 1;
    static const uint32_t lineWordIdx = 2;
    
    const uint64_t memSizeGB;
    const uint64_t memSize;
    const uint64_t numHuges;
    const uint64_t numPages;
    
    const double k_alloc;
    const double u_alloc;
    const double p_dup_huge;
    const double p_dup_page;
    const double p_dup_line;
    
    set< uint64_t > globHuges;
    set< uint64_t > globPages;
    set< uint64_t > globLines;

    set< uint64_t > kernHuges;
    set< uint64_t > kernPages;
    set< uint64_t > kernLines;

    set< uint64_t > userHuges;
    set< uint64_t > userPages;
    set< uint64_t > userLines;
    
    RandomNumberGenerator rng;
    PowerLaw hugeDist;
    PowerLaw pageDist;
    PowerLaw lineDist;
    
    ofstream imgFile;
    ofstream refFile;

    void ZeroLines( const uint32_t numZeros );
    void PagesLoop( void );
    void LinesLoop( void );
    void DuplicatePage( void );
    void DuplicateHuge( void );
    void DuplicateLine( void );
    void NonRandomHuge( const uint64_t idx );
    void NonRandomPage( const uint64_t idx );
    void NonRandomLine( const uint64_t idx );
    void RandomLine( void );

    void WritePageRefCounts( const uint32_t numPages );
    
    uint64_t globZeroHuges;
    uint64_t globZeroPages;
    uint64_t globZeroLines;
    uint64_t globUniqHuges;
    uint64_t globUniqPages;
    uint64_t globUniqLines;
    
    uint64_t kernZeroHuges;
    uint64_t kernZeroPages;
    uint64_t kernZeroLines;
    uint64_t kernUniqHuges;
    uint64_t kernUniqPages;
    uint64_t kernUniqLines;

    uint64_t userZeroHuges;
    uint64_t userZeroPages;
    uint64_t userZeroLines;
    uint64_t userUniqHuges;
    uint64_t userUniqPages;
    uint64_t userUniqLines;
    
    line_t zeroLine;
    
    bool kernel;
    bool user;
public:
    ContentGenerator( const uint32_t _gb, const double _k, const double _u, const double _h, const double _p, const double _l ) :
    memSizeGB( _gb ),
    memSize( 1024ULL * 1024ULL * 1024ULL * memSizeGB ),
    numHuges( memSize / hugeSize ),
    numPages( memSize / pageSize ),
    k_alloc( _k ),
    u_alloc( _u ),
    p_dup_huge( _h ),
    p_dup_page( _p ),
    p_dup_line( _l ),
    hugeDist( 2.0 ),
    pageDist( 0.5 ),
    lineDist( 0.5 ) {
        imgFile.open( "memimage.bin.0000", ios::binary );
        refFile.open( "refcount.bin.0000", ios::binary );

        globZeroHuges = 0ULL;
        globZeroPages = 0ULL;
        globZeroLines = 0ULL;
        globUniqHuges = 0ULL;
        globUniqPages = 0ULL;
        globUniqLines = 0ULL;

        kernZeroHuges = 0ULL;
        kernZeroPages = 0ULL;
        kernZeroLines = 0ULL;
        kernUniqHuges = 0ULL;
        kernUniqPages = 0ULL;
        kernUniqLines = 0ULL;

        userZeroHuges = 0ULL;
        userZeroPages = 0ULL;
        userZeroLines = 0ULL;
        userUniqHuges = 0ULL;
        userUniqPages = 0ULL;
        userUniqLines = 0ULL;

        zeroLine.SetToZero();
        kernel = false;
        user   = false;
    }
    ~ContentGenerator( void ) {
        imgFile.close();
        refFile.close();
    }
    void Generate( void );
    void PrintStats( void ) const;
};
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ContentGenerator::PrintStats( void ) const {
    cout << "all memory:" << endl;
    cout << "... huges:" << endl;
    cout << "... ... zero:   " << PrettyInt( globZeroHuges, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( globUniqHuges, 16 ) << endl;
    cout << "... pages:" << endl;
    cout << "... ... zero:   " << PrettyInt( globZeroPages, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( globUniqPages, 16 ) << endl;
    cout << "... lines:" << endl;
    cout << "... ... zero:   " << PrettyInt( globZeroLines, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( globUniqLines, 16 ) << endl;

    cout << "kernel memory:" << endl;
    cout << "... huges:" << endl;
    cout << "... ... zero:   " << PrettyInt( kernZeroHuges, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( kernUniqHuges, 16 ) << endl;
    cout << "... pages:" << endl;
    cout << "... ... zero:   " << PrettyInt( kernZeroPages, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( kernUniqPages, 16 ) << endl;
    cout << "... lines:" << endl;
    cout << "... ... zero:   " << PrettyInt( kernZeroLines, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( kernUniqLines, 16 ) << endl;

    cout << "user memory:" << endl;
    cout << "... huges:" << endl;
    cout << "... ... zero:   " << PrettyInt( userZeroHuges, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( userUniqHuges, 16 ) << endl;
    cout << "... pages:" << endl;
    cout << "... ... zero:   " << PrettyInt( userZeroPages, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( userUniqPages, 16 ) << endl;
    cout << "... lines:" << endl;
    cout << "... ... zero:   " << PrettyInt( userZeroLines, 16 ) << endl;
    cout << "... ... unique: " << PrettyInt( userUniqLines, 16 ) << endl;
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ContentGenerator::ZeroLines( const uint32_t numZeros ) {
    for( uint32_t i = 0; i < numZeros; i++ ) {
        imgFile.write( zeroLine.RawCharPtr(), lineSize );
    }
}
void ContentGenerator::RandomLine( void ) {

    globUniqLines++;
    if( kernel ) { kernUniqLines++; }
    if( user   ) { userUniqLines++; }

    line_t line;
    for( uint32_t i = 0; i < wordsPerLine; i++ ) {
        line.SetWord( rng.RandU64(), i );
    }

    imgFile.write( line.RawCharPtr(), lineSize );
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ContentGenerator::WritePageRefCounts( const uint32_t numPages ) {

    const double pk = rng.RandDouble();
    const double pu = rng.RandDouble();

    kernel = pk < k_alloc ? true : false;
    user   = pu < u_alloc ? true : false;
    user   = kernel ? user : false;

    const uint32_t krc = kernel ? 1 : 0;
    const uint32_t urc = user   ? 0 : 0xFFFFFFFF;

    for( uint32_t pageIdx = 0; pageIdx < numPages; pageIdx++ ) {

        refFile.write( (char*) &krc, sizeof( uint32_t ) );
        refFile.write( (char*) &urc, sizeof( uint32_t ) );
    }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ContentGenerator::NonRandomHuge( const uint64_t idx ) {

    if( !globHuges.count( idx ) ) {
        globUniqLines += linesPerHuge;
        globUniqPages += pagesPerHuge;
        globUniqHuges++;
        globHuges.insert( idx );
    }
    if( kernel ) {
        if( !kernHuges.count( idx ) ) {
            kernUniqLines += linesPerHuge;
            kernUniqPages += pagesPerHuge;
            kernUniqHuges++;
            kernHuges.insert( idx );
        }
    }
    if( user ) {
        if( !userHuges.count( idx ) ) {
            userUniqLines += linesPerHuge;
            userUniqPages += pagesPerHuge;
            userUniqHuges++;
            userHuges.insert( idx );
        }
    }
    
    line_t line;
    line.SetToZero();

    uint64_t word = idx * linesPerHuge;
    line.SetWord( word, hugeWordIdx );
    
    for( uint32_t lineIdx = 0; lineIdx < linesPerHuge; lineIdx++ ) {
        
        imgFile.write( line.RawCharPtr(), lineSize );
        word++;
        line.SetWord( word, hugeWordIdx );
    }
}
void ContentGenerator::NonRandomPage( const uint64_t idx ) {

    if( !globPages.count( idx ) ) {
        globUniqLines += linesPerPage;
        globUniqPages++;
        globPages.insert( idx );
    }
    if( kernel ) {
        if( !kernPages.count( idx ) ) {
            kernUniqLines += linesPerPage;
            kernUniqPages++;
            kernPages.insert( idx );
        }
    }
    if( user ) {
        if( !userPages.count( idx ) ) {
            userUniqLines += linesPerPage;
            userUniqPages++;
            userPages.insert( idx );
        }
    }
    
    line_t line;
    line.SetToZero();

    uint64_t word = idx * linesPerPage;
    line.SetWord( word, pageWordIdx );
    
    for( uint32_t lineIdx = 0; lineIdx < linesPerPage; lineIdx++ ) {
        
        imgFile.write( line.RawCharPtr(), lineSize );
        word++;
        line.SetWord( word, pageWordIdx );
    }
}
void ContentGenerator::NonRandomLine( const uint64_t idx ) {

    if( !globLines.count( idx ) ) {
        globUniqLines++;
        globLines.insert( idx );
    }
    if( kernel ) {
        if( !kernLines.count( idx ) ) {
            kernUniqLines++;
            kernLines.insert( idx );
        }
    }
    if( user ) {
        if( !userLines.count( idx ) ) {
            userUniqLines++;
            userLines.insert( idx );
        }
    }
    
    line_t line;
    line.SetToZero();

    line.SetWord( idx, lineWordIdx );
    imgFile.write( line.RawCharPtr(), lineSize );
    
}
void ContentGenerator::DuplicateHuge( void ) {
    const uint64_t idx = hugeDist.IntSample() - 1ULL;
    
    WritePageRefCounts( pagesPerHuge );
    
    if( idx == 0ULL ) {
        ZeroLines( linesPerHuge );
        globZeroHuges++;
        globZeroPages += pagesPerHuge;
        globZeroLines += linesPerHuge;
        if( kernel ) {
            kernZeroHuges++;
            kernZeroPages += pagesPerHuge;
            kernZeroLines += linesPerHuge;
        }
        if( user ) {
            userZeroHuges++;
            userZeroPages += pagesPerHuge;
            userZeroLines += linesPerHuge;
        }
    }
    else {
        NonRandomHuge( idx );
    }
}
void ContentGenerator::DuplicatePage( void ) {
    const uint64_t idx = pageDist.IntSample() - 1ULL;
    
    if( idx == 0ULL ) {
        ZeroLines( linesPerPage );
        globZeroPages++;
        globZeroLines += linesPerPage;
        if( kernel ) {
            kernZeroPages++;
            kernZeroLines += linesPerPage;
        }
        if( user ) {
            userZeroPages++;
            userZeroLines += linesPerPage;
        }
    }
    else {
        NonRandomPage( idx );
    }
}
void ContentGenerator::DuplicateLine( void ) {
    const uint64_t idx = pageDist.IntSample() - 1ULL;

     if( idx == 0ULL ) {
         ZeroLines( 1 );
         globZeroLines++;
         if( kernel ) { kernZeroLines++; }
         if( user )   { userZeroLines++; }
     }
     else {
         NonRandomLine( idx );
     }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
void ContentGenerator::LinesLoop( void ) {
    for( uint32_t lineIdx = 0; lineIdx < linesPerPage; lineIdx++ ) {
        
        if( lineIdx == 0 ) {
            globUniqPages++;
            if( kernel ) { kernUniqPages++; }
            if( user )   { userUniqPages++; }
        }
        
        const double p = rng.RandDouble();
        if( lineIdx && p < p_dup_line ) {
            DuplicateLine();
        }
        else {
            RandomLine();
        }
    }
}
void ContentGenerator::PagesLoop( void ) {
    for( uint32_t pageIdx = 0; pageIdx < pagesPerHuge; pageIdx++ ) {

        WritePageRefCounts( 1 );
        if( pageIdx == 0 ) {
            globUniqHuges++;
            if( kernel ) { kernUniqHuges++; }
            if( user )   { userUniqHuges++; }
        }

        const double p = rng.RandDouble();
        if( pageIdx && p < p_dup_page ) {
            DuplicatePage();
        }
        else {
            LinesLoop();
        }
    }
}
void ContentGenerator::Generate( void ) {
    for( uint32_t hugeIdx = 0; hugeIdx < numHuges; hugeIdx++ ) {
        
        const double p = rng.RandDouble();
        if( p < p_dup_huge ) {
            DuplicateHuge();
        }
        else {
            PagesLoop();
        }        
    }
}
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
int main( int argc, char ** argv ) {
    if( argc != 7 ) {
        cout << "usage: " << argv[0] << " <memsize_gb> <k_alloc> <u_alloc> <p_dup_huge> <p_dup_page> <p_dup_line>" << endl;
        cout << "... memsize_gb: memory size in GB, must be divisible by 4" << endl;
        cout << "... k_alloc: percent of memory allocated" << endl;
        cout << "... u_alloc: percent of kernel memory allocated to user space" << endl;
        cout << "... p_dup_huge: probability of a duplicated 2MB page" << endl;
        cout << "... p_dup_page: probability of a duplicated 4kB page, if 2MB region is not duplicated" << endl;
        cout << "... p_dup_line: probability of a dupcliated 64B line, if 4kB region is not duplicated" << endl;
        exit( -1 );
    }
    const uint32_t memsize_gb = atoi( argv[1] );
    const double k_alloc = atof( argv[2] );
    const double u_alloc = atof( argv[3] );
    const double p_dup_huge = atof( argv[4] );
    const double p_dup_page = atof( argv[5] );
    const double p_dup_line = atof( argv[6] );
    assert( memsize_gb % 4 == 0 );
    assert( memsize_gb < 129 );
    assert( k_alloc >= 0.0 && k_alloc <= 1.0 );
    assert( u_alloc >= 0.0 && u_alloc <= 1.0 );
    assert( p_dup_huge >= 0.0 && p_dup_huge <= 1.0 );
    assert( p_dup_page >= 0.0 && p_dup_page <= 1.0 );
    assert( p_dup_line >= 0.0 && p_dup_line <= 1.0 );
    
    ContentGenerator gen( memsize_gb, k_alloc, u_alloc, p_dup_huge, p_dup_page, p_dup_line );
    gen.Generate();
    gen.PrintStats();
    
    return 0; 
}




